import argparse
import os
import numpy as np
# 导入自定义工具
from utils.h5data import read_h5_data, save_h5_data
from utils.phy import *
from myglobal import *
def fit_slope_only(x, y):
    """
    只拟合斜率k，固定截距b=0，即 y = k*x
    使用最小二乘法: k = (x^T * y) / (x^T * x)
    同时计算k的标准误差
    """
    # 计算斜率 k
    k = np.dot(x, y) / np.dot(x, x)
    
    # 计算残差
    residuals = y - k * x
    n = len(x)
    
    # 计算均方误差
    mse = np.sum(residuals**2) / (n - 1)  # 自由度为 n-1 (因为我们估计了1个参数)
    
    # 计算k的方差和标准误差
    k_var = mse / np.dot(x, x)
    k_err = np.sqrt(k_var)
    
    
    return k, k_err, mse

def get_dvv_data_from_7file(INPUT_FILE, STATS=None, v_ref=3000):
    '''
    测量dv/v
    '''
    # 获取所有台站名称
    if STATS is None:
        all_groups = read_h5_data(INPUT_FILE, ['all_groups'], group_name='metadata')[0]
        all_groups = [i.decode() for i in all_groups]
    else:
        all_groups = STATS
    
    data_all = {}
    
    for j, name in enumerate(all_groups):
        delta_t, te, T_MAX, r = read_h5_data(INPUT_FILE, 
                    ['delta_t','te','T_MAX','r'],
                    group_name=name)
        # if abs(r) >= 3000 or abs(r)<100:  # 只选择距离小于1500的台站
        #     continue
        # if abs(r)<=10:
        #     continue
        if name not in ['P300']:
            continue
        print(j, name,r)
        data_all[name] = [delta_t, te, T_MAX, r]

    if len(data_all.keys()) < 1:
        print('total data:', len(data_all.keys()))
        raise ValueError('not enough data')
    else:
        print('total data:', len(data_all.keys()))

    st_valid = list(data_all.keys())
    ne = len(data_all[st_valid[0]][0])  # 事件数量
    ns = len(st_valid)  # 台站数量
    
    dvv = np.zeros(ne)
    dvv_err = np.zeros(ne)
    dt_err = np.zeros(ne)
    fit_quality = np.zeros(ne)  # 存储拟合质量(R²值)

    DT_ALL = [data_all[name][0] for name in st_valid]
    dt_var = np.var(DT_ALL)
    
    # 对每个事件进行处理
    for it in range(ne):
        delta_t_i = []
        x = []
        
        # 收集该事件下所有台站的数据
        for j, name in enumerate(st_valid):
            delta_t_val = data_all[name][0][it]
            xj = data_all[name][3]
            # 只使用有效的数据点（非零且有限值）
            if np.isfinite(delta_t_val) and delta_t_val != 0:
                delta_t_i.append(delta_t_val)
                x.append(xj)  # 距离
        
        # 如果有足够的有效数据点，则进行线性拟合
        if len(x) >= 2:
            x_valid = np.array(x)
            delta_t_valid = np.array(delta_t_i)
            
            k,k_err, mse = fit_slope_only(x_valid, delta_t_valid)

            dt_err[it] = 0
            dvv[it] = -k * v_ref
            dvv_err[it] = k_err * v_ref
            # fit_quality[it] = np.sqrt(mse/dt_var)
            fit_quality[it] = np.sqrt(mse)
        elif len(x) == 1:
            dvv[it] = -delta_t_i[0]/x[0] * v_ref
            dvv_err[it] = 0
            dt_err[it] = 0
            fit_quality[it] = 0
        else:
            dvv[it] = 0
            dvv_err[it] = 0
            dt_err[it] = 0
            fit_quality[it] = 0
    
    return te, dvv, dvv_err, dt_err,fit_quality

def plot_dvv_curve(te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''
    from pylab import figure, rcParams, plt
    
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    
    fig = figure(figsize=(8, 2.5), dpi=500)
    
    # 主图：dv/v vs 时间, 原始单位： ms/s，所以是‰
    ax1 = plt.subplot(1, 1, 1)

    print(te_bjt[0],te_bjt[-1])
    print(dvv)
    plt.errorbar(te_bjt, dvv, yerr=dvv_err*3, 
                 fmt='.-', linewidth=0.7,color='tab:red', markersize=3, markeredgecolor='none',
                ecolor='none',elinewidth=0.6, capsize=1, capthick=0.6, barsabove=True
                     )
    plt.fill_between(te_bjt, dvv-dvv_err*3, dvv+dvv_err*3, alpha=0.4, color='tab:red')

    
    plt.xlabel('day (BJT)')
    plt.ylabel('dv/v(%)',color='tab:red')
    plt.title(title)
    plt.grid(True, alpha=0.3)
    
    plt.ylim([-np.std(dvv)*4,np.std(dvv)*4])

    ax2 = plt.twinx(ax1)
    plt.plot(te_bjt, fit_quality, '-', color='k',linewidth=0.6)
    plt.ylabel('rms(ms)', color='k')
    plt.ylim(0, np.std(fit_quality)*20)
    output_file = output_file.replace('.png', '.rms.png')
    # plt.yticks([0,50,100])

    plt.tight_layout()
    
    if output_file:
        plt.savefig(output_file, dpi=600, bbox_inches='tight')
        print(f"Figure saved to {output_file}")
    
    return fig


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot traceSeq h5 data')
    parser.add_argument('-input', default='', help='input h5 file')
    parser.add_argument('-figroot', default='figures/7.dt.change.figures', help='root to save figs')
    parser.add_argument('-vref', type=float, default=3000, help='reference velocity (m/s)')
    # parser.add_argument('-figroot', default='figures/debug', help='root to save figs')
    args = parser.parse_args()
    INPUT_FILE = args.input
    FIG_ROOT = args.figroot
    V_REF = args.vref

    datasets, args_in_file = read_h5_data(INPUT_FILE, 
                                         ['all_groups','MARKER','EMARKER','DATE'],
                                         group_name='metadata',
                                         read_attrs=True)
    all_groups = [i.decode() for i in datasets[0]]
    MARKER = datasets[1].decode() 
    EMARKER = datasets[2].decode() 
    DATE = datasets[3].decode() 

    FIG_ROOT = f'{FIG_ROOT}/7B.{DATE}.{EMARKER}'
    if not os.path.exists(FIG_ROOT):
        os.makedirs(FIG_ROOT)

    te, dvv, dvv_err, dt_err, fit_quality = get_dvv_data_from_7file(INPUT_FILE, v_ref=V_REF)
    te_bjt = te/3600/24 + 8/24

    # 保存数据，与INPUT_FILE同一个文件夹下，但是名称另取
    folder = os.path.dirname(INPUT_FILE)
    OUTFILE = f'{folder}/dvv.{MARKER}.V{V_REF:06.1f}.h5'
    save_h5_data(OUTFILE, 
                 data_in_dict={'te_bjt':te_bjt, 'dvv':dvv, 'dvv_err':dvv_err, 'dt_err':dt_err, 'fit_quality':fit_quality,
                               'all_groups':all_groups,'MARKER':MARKER, 'EMARKER':EMARKER, 'DATE':DATE},)

    mask = np.where(np.all([te_bjt>=0,te_bjt<=180],axis=0))
    # mask = np.where(te_bjt<=10000)
    dvv     = -dvv[mask]/10 # 0.001转换为百分比,并校正X的符号
    dvv_err = dvv_err[mask]/10
    dt_err  = dt_err[mask]
    te_bjt  = te_bjt[mask]
    fit_quality = fit_quality[mask]
    print(te_bjt)

    # 绘制dv/v曲线
    dvv_fig = plot_dvv_curve(te_bjt, dvv, dvv_err, fit_quality,
                             title=f'{MARKER} DV/V vs Time',
                            output_file=f'{FIG_ROOT}/dvv_vs_time.{MARKER}.V{V_REF:06.1f}.png')